import argparse
import pandas as pd 
import matplotlib.pyplot as plt
import numpy as np


def plot_pos_plane(df: pd.DataFrame, keypoints:list[int], cmap='winter', 
                   video_path=None, video_sec=0):

    if video_path is not None:
        import cv2
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        fp_idx = int(round(fps * video_sec))
        cap.set(cv2.CAP_PROP_POS_FRAMES, fp_idx)
        ret, frame = cap.read()

        if ret:
            plt.imshow(frame)
        else:
            print('Unable to load BG image')
        cap.release()

    for k in keypoints:
        dfk = df[df['label'] == k]
        plt.scatter(dfk['x'], dfk['y'], c=dfk['frame'], cmap=cmap, label=str(k))

    plt.gca().yaxis.set_inverted(True)
    plt.gca().set_aspect('equal')
    # plt.colorbar()
    plt.show()


def plot_time_dependent(df:pd.DataFrame, keypoints: list[int], title=None):
    for k in keypoints:
        dfk = df[df['label'] == k]
        plt.plot(dfk['frame'], dfk['x'], '-', label=f'{k}x')
        plt.plot(dfk['frame'], dfk['y'], '--', label=f'{k}y')
    plt.xlabel(f'frame')
    plt.ylabel(f'pos (pixel)')
    plt.legend()
    if title is not None:
        plt.title(title)
    plt.show()


def plot_td_angle(df:pd.DataFrame, center_point_id, left_point_id, right_point_id, label):
    frames = []
    angles = []
    
    df = df.set_index('frame')
    frame_indexes = df.index.unique()

    for f in frame_indexes:
        dff = df.loc[f, :]
        dff = dff.set_index('label')
        if center_point_id in dff.index and left_point_id in dff.index and right_point_id in dff.index:
            # This is a valid frame
            vc = np.array(dff.loc[center_point_id, ['x', 'y']])
            vl = np.array(dff.loc[left_point_id, ['x', 'y']])
            vr = np.array(dff.loc[right_point_id, ['x', 'y']])
            
            v1 = vl - vc
            v2 = vr - vc

            angle_rad = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
            angle_deg = angle_rad * 180 / np.pi
            frames.append(f)
            angles.append(angle_deg)

    plt.plot(frames, angles, label=label)
    plt.xlabel('frame')
    plt.ylabel('angle (degree)')

def plot_td_elbow_angle(df:pd.DataFrame, side:str, title=None):
    if side == 'left' or side == 'both':
        plot_td_angle(df, 7, 5, 9, 'left_elbow')
    if side == 'right' or side == 'both':
        plot_td_angle(df, 8, 6, 10, 'right_elbow')
    if side not in {'left', 'right', 'both'}:
        raise ValueError(f'Invalid side {side=}')
    
    if title is not None:
        plt.title(f'{title}  elbow angle')
    else:
        plt.title('elbow angle')
    
    plt.xlabel('frame')
    plt.ylabel('elbow angle (degree)')
    plt.show()


def plot_td_argument_angle(df:pd.DataFrame, start_point:int, end_point:int, label=None, title=None):
    """
    2025.01.10  The azimuth angle (argument angle) of the vector defined by start and end point (keypoint id).
    """
    frames = []
    angles = []
    
    df = df.set_index('frame')
    frame_indexes = df.index.unique()

    for f in frame_indexes:
        dff = df.loc[f, :]
        dff = dff.set_index('label')
        if start_point in dff.index and end_point in dff.index:
            # This is a valid frame
            vs = np.array(dff.loc[start_point, ['x', 'y']])
            ve = np.array(dff.loc[end_point, ['x', 'y']])
            
            dv = ve - vs

            angle_rad = np.arctan2(-dv[1], dv[0])
            angle_deg = angle_rad * 180 / np.pi
            frames.append(f)
            angles.append(angle_deg)
    plt.plot(frames, angles, label=label)
    plt.ylabel('argument angle (degree)')
    plt.xlabel('frame')


def plot_arrow_arg_angle(df, strech_side, title=None):
    if strech_side == 'left':
        plot_td_argument_angle(df, 9, 10, 'arrow_arg_left', title)
    else:
        plot_td_argument_angle(df, 10, 9, 'arrow_ang_right', title)
    
    plt.axhline(0, ls=':', c='k')
    plt.ylabel('arrow argument angle (degree)')
    if title is not None:
        plt.title(f'{title}  arrow argument angle')
    else:
        plt.title('arrow argument angle')
    plt.show()
    


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('source_file', default='keypoints.csv')
    parser.add_argument('task', help='task type, including: plane, time')
    parser.add_argument('--confident', '-c', type=float, default=0.8)
    parser.add_argument('--cmap', default='winter')
    parser.add_argument('--keypoints', '-k', nargs='*', required=True, type=int)
    parser.add_argument('--background-video', '-b', default=None)
    parser.add_argument('--background-timepoint', '-t', type=float, default=0.0,
                        help='Timepoint (second) for the bg figure in the given video')
    parser.add_argument('--title', default=None)

    p = parser.parse_args()

    df = pd.read_csv(p.source_file)
    df = df[df['confident'] >= p.confident]

    match (p.task):
        case 'plane':
            plot_pos_plane(df, p.keypoints, p.cmap, p.background_video, p.background_timepoint)
        case 'time':
            plot_time_dependent(df, p.keypoints, p.title)
        case _:
            print(f'Invalid task: {p.task}')

